package com.jin.webvideo.base.dc;

import javassist.bytecode.ClassFile;
import javassist.compiler.CompileError;

import javax.tools.*;
import java.io.*;
import java.util.*;

/**
 * DCClassLoader创建者
 * @author biao532
 * @date 2023/12/02 19:10
 **/
public class DCClassLoaderCreator {
    public static DCClassLoaderCreator creator() {
        return new DCClassLoaderCreator();
    }

    private List<MemoryFileManager.OutputMemoryJavaFileObject> list = new ArrayList<>(0);
    private DCClassLoaderCreator(){
    }

    public ClassLoader createClassLoader(ClassLoader parent) {
        // 每次new, 防止类缓存
        DynamicCompilationClassLoader f = new DynamicCompilationClassLoader("test", parent);
        for (MemoryFileManager.OutputMemoryJavaFileObject outputMemoryJavaFileObject : list) {
            f.register(outputMemoryJavaFileObject);
        }
        return f;
    }

    public DCClassLoaderCreator addByteCode(byte[] byteCode) throws IOException {
        try(InputStream inputStream = new ByteArrayInputStream(byteCode)) {
            ClassFile file = new ClassFile(new DataInputStream(inputStream));
            MemoryFileManager.OutputMemoryJavaFileObject o = new MemoryFileManager.OutputMemoryJavaFileObject(file.getName(), JavaFileObject.Kind.CLASS);
            o.openOutputStream().write(byteCode);
            return this;
        }
    }

    public DCClassLoaderCreator addClassFile(InputStream inputStream) throws IOException {
        try(ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
            inputStream.transferTo(stream);
            byte[] bs = stream.toByteArray();
            return addByteCode(bs);
        }
    }

    public DCClassLoaderCreator compile(List<String> args, String... names) throws CompileError {
        JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
        StandardJavaFileManager standardFileManager = compiler.getStandardFileManager(null, null, null);
        MemoryFileManager fileManager = new MemoryFileManager(standardFileManager);

        fileManager.registerClassFileCreationListener(new MemoryFileManager.ClassFileCreationListener() {
            @Override
            public void newClassFile(MemoryFileManager.OutputMemoryJavaFileObject jfo, JavaFileManager.Location location, String className, JavaFileObject.Kind kind, FileObject sibling) {
                if (kind== JavaFileObject.Kind.CLASS) {
                    list.add(jfo);
                }
            }
        });

        Iterable<? extends JavaFileObject> fileObjects = standardFileManager.getJavaFileObjects(names);
        // 使用stringWriter来收集错误。
        StringWriter errorStringWriter = new StringWriter();
        JavaCompiler.CompilationTask task = compiler.getTask(errorStringWriter, fileManager, (e) -> {
            if (e.getKind() == Diagnostic.Kind.ERROR) {
                errorStringWriter.append(e.toString());
            }

        }, args, null, fileObjects);
        if (!task.call()) {
            throw new CompileError(errorStringWriter.toString());
            //return;
        }
        return this;
    }
}
